from collections import defaultdict

from app.utils.cv_tools import show_image_with_roi, crop_image_with_roi
from tests.base_test import base_test_case
import shutil
from pathlib import Path

import cv2
import numpy as np
from paddleocr import LayoutDetection

from app.utils.pdf_tools import PDFExtractor
import json

TEST_DATA_DIR = base_test_case.test_data_dir
SUB_DIR = TEST_DATA_DIR / "1715339805571"


# print(os.environ.get("PADDLE_PDX_CACHE_HOME"))


# os.environ['PADDLE_OCR_BASE_DIR'] = '/data/repo/PADDLE_OCR_BASE_DIR' # 模型缓存目录未生效
def extract_pdf_pages_as_images(pdf_path, page_range=None, rotation=0, zoom=2.0):
    """
    从PDF文件中提取每页作为图像

    Args:
        pdf_path (str): PDF文件路径
        page_range (tuple): 页面范围，例如 (0, 3) 表示处理第1到第3页
        rotation (int): 旋转角度，默认为0
        zoom (float): 缩放比例，默认为2.0
                   基于96 DPI基准，换算关系为：DPI = 96 * zoom
                   常用设置：
                   - 屏幕显示：zoom=1.0-2.0（96-192 DPI）
                   - OCR识别：zoom=2.0-3.0（192-288 DPI）
                   - 打印/出版：zoom=3.125（300 DPI）

    Yields:
        tuple: (page_index, numpy.ndarray) 页面索引和每页的图像数据
    """
    from paddlex.utils.deps import is_dep_available
    if is_dep_available("pypdfium2"):
        import pypdfium2 as pdfium
    else:
        print("pypdfium2 is not available")
        # logger.error("pypdfium2 is not available")
        return

    # 打开PDF文档
    doc = pdfium.PdfDocument(pdf_path)
    try:
        page_count = len(doc)

        # 确定要处理的页面范围
        start_page, end_page = 0, page_count
        if page_range:
            start_page, end_page = page_range
            # 确保页码范围有效
            start_page = max(0, start_page)
            end_page = min(page_count, end_page)

        # 遍历指定范围内的页面
        for page_index in range(start_page, end_page):
            # 获取指定页面
            page = doc[page_index]
            # 将页面渲染为图像
            image = page.render(rotation=rotation, scale=zoom).to_pil()
            # 转换为RGB模式确保颜色一致性
            image = image.convert("RGB")
            # 转换为OpenCV格式（BGR）
            img_cv = np.array(image)
            img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR)
            # 返回页面索引和图像数据
            yield page_index, img_cv
    finally:
        # 确保在使用完毕后关闭PDF文档
        doc.close()


def layout_detection_demo1():
    # 初始化
    init_kwargs = {
        'device': 'cpu', 'cpu_threads': 10,
        # "threshold": 0.5,
        # "layout_nms": True,
        # "model_dir": "", # /home/tomcat/.paddlex/official_models/PP-DocLayout_plus-L
        # "model_name": "layout",
    }
    # 修复：直接使用图片路径而不是包含在列表中
    image_paths = [
        # TEST_DATA_DIR / "1715339805571.pdf",
        TEST_DATA_DIR / "25-注会-轻1-财务成本管理[上册](第3章).pdf",
        # TEST_DATA_DIR / "doc_with_formula.png",
        # SUB_DIR / "1715339805571_001.jpg",
        # SUB_DIR / "1715339805571_002.jpg",
        # SUB_DIR / "1715339805571_003.jpg",
        # SUB_DIR / "1715339805571_004.jpg",
        # SUB_DIR / "1715339805571_005.jpg",
    ]

    # 逐个处理图像而不是传递整个列表
    for image_path in image_paths:
        # 如果是PDF文件，则提取页面并直接传递numpy数组
        if str(image_path).endswith('.pdf'):
            # 指定要处理的页面范围 (从0开始索引，例如(0, 3)表示处理第1到第3页)
            page_range = (0, 1)  # 可根据需要修改此范围
            page_images = []
            page_indices = []

            # 提取指定范围的PDF页面
            for page_index, page_image in extract_pdf_pages_as_images(str(image_path), page_range=page_range, zoom=1.0):
                page_images.append(page_image)
                page_indices.append(page_index)

            layout_detection_predictor = LayoutDetection(**init_kwargs)
            # 推理
            result = layout_detection_predictor.predict(page_images, batch_size=1)

            # 创建PDF输出目录
            pdf_name = image_path.stem
            output_dir = Path("./output") / pdf_name
            if output_dir.exists():
                shutil.rmtree(output_dir)
            output_dir.mkdir(parents=True, exist_ok=True)

            for i, res in enumerate(result):
                # 添加按cls_id升序排序的代码
                res.get("boxes", []).sort(key=lambda x: x['cls_id'])

                # 按页码保存图像和JSON文件
                page_filename = f"page_{page_indices[i]}"
                res.save_to_img(save_path=str(output_dir / f"{page_filename}.jpg"))
                res.save_to_json(save_path=str(output_dir / f"{page_filename}.json"))

                # 查找所有label为"image"的元素坐标
                image_boxes = []
                for box in res.get("boxes", []):
                    if box.get("label") == "image":
                        image_boxes.append(box)
                # 如果找到了image标签的坐标，则使用它们进行裁剪
                for idx, box in enumerate(image_boxes):
                    coordinate = box.get("coordinate")
                    # 为每个图像裁剪结果生成不同的文件名
                    crop_filename = f"{page_filename}_crop_{idx}.jpg" if len(image_boxes) > 1 else f"{page_filename}_image_crop.jpg"

                    crop_image_with_roi(
                        res['input_img'],
                        roi=coordinate,
                        output_path=str(output_dir / crop_filename),
                        expand_roi=1  # 外扩1个像素确保边界完整
                    )
                    show_image_with_roi(res['input_img'], roi=coordinate, is_show=True)

                table_boxes = []
                for box in res.get("boxes", []):
                    if box.get("label") == "table":
                        table_boxes.append(box)
                for idx, box in enumerate(table_boxes):
                    coordinate = box.get("coordinate")
                    # 为每个图像裁剪结果生成不同的文件名
                    crop_filename = f"{page_filename}_crop_{idx}.jpg" if len(table_boxes) > 1 else f"{page_filename}_table_crop.jpg"

                    crop_image_with_roi(
                        res['input_img'],
                        roi=coordinate,
                        output_path=str(output_dir / crop_filename),
                        expand_roi=1  # 外扩1个像素确保边界完整
                    )
                    show_image_with_roi(res['input_img'], roi=coordinate, is_show=True)
        else:
            # 将Path对象转换为字符串
            result = layout_detection_predictor.predict(str(image_path))
            for res in result:
                # res.print()
                # res.print(json_format=True)
                res.save_to_img(save_path="./output/")
                res.save_to_json(save_path="./output/")


def pdf_extractor_demo1():
    """
    测试PDF提取器功能
    """
    # 初始化提取器
    extractor = PDFExtractor()

    # 检查是否可用
    if not extractor.is_available():
        print("pypdfium2不可用，请先安装: pip install pypdfium2")
        return

    # PDF文件路径
    pdf_path = TEST_DATA_DIR / "1715339805571.pdf"

    if not pdf_path.exists():
        print(f"PDF文件不存在: {pdf_path}")
        return

    print(f"处理PDF文件: {pdf_path}")

    # 获取页数
    page_count = extractor.get_page_count(str(pdf_path))
    print(f"PDF总页数: {page_count}")

    # 提取第一页信息
    if page_count > 0:
        page_info = extractor.extract_page_info(str(pdf_path), 0)
        print(f"第一页信息: {page_info}")

    # 提取所有页面为图像
    print("提取页面为图像...")
    image_files = extractor.extract_to_images(
        str(pdf_path),
        "./output/pdf_images",
        image_format="jpg",
        zoom=2.0
    )
    print(f"已保存 {len(image_files)} 个图像文件")

    # 提取所有页面为文本文件
    print("提取页面为文本...")
    text_files = extractor.extract_to_text_files(
        str(pdf_path),
        "./output/pdf_texts"
    )
    print(f"已保存 {len(text_files)} 个文本文件")

    # 提取为JSON格式
    print("提取为JSON格式...")
    json_file = extractor.extract_to_json(
        str(pdf_path),
        "./output/pdf_content.json"
    )
    print(f"已保存JSON文件: {json_file}")


if __name__ == '__main__':
    layout_detection_demo1()
    # test_pdf_extractor()

    pass
